package UnionFind;

public class WeightedQuickUnionUF {
    private int[] ids;
    private int[] weights;
    private int count;
    public WeightedQuickUnionUF(int N) {
        count = N;
        this.ids = new int[N];
        this.weights = new int[N];
        for (int i = 0; i < N; i++) {
            ids[i] = i;
            weights[i] = 1;
        }
    }

    public int find(int q) {
        while (q != ids[q]) {
            q = ids[q];
        }
        return q;
    }

    public int count() {
        return count;
    }

    public boolean connected(int p, int q) {
        return find(p) == find(q);
    }

    public void union(int p, int q) {
        int pid = find(p);
        int qid = find(q);
        if (pid == qid) {
            return;
        }

        if (weights[pid] >= weights[qid]) {
            ids[qid] = pid;
            weights[qid] += weights[pid];
        } else {
            ids[pid] = qid;
            weights[pid] = weights[qid];
        }
        count--;
    }
}
